# -*- coding: utf-8 -*-
"""
Created on Sun Nov 14 19:10:04 2021

@author: zhiyinan
"""
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.font_manager as font_manager
from casadi import *

Nu = 1
Ny = 1
past = 20
future = 1

y_index = 19

Delta = 2*60*60
Nt = 40*60*60
Nsim = 5*24*60*60

csfont = {'fontname':'Times New Roman'}
font = font_manager.FontProperties(family='Times New Roman',
                                    # weight='bold',
                                    style='normal', size=15)

font_l = font_manager.FontProperties(family='Times New Roman',
                                   # weight='bold',
                                    style='normal', size=16)

Precip = np.loadtxt("ZEMPC/Precip.txt")

# %% 2% Porcess Noise - ODE, LSTM & 2-layered
u_ODE = np.loadtxt("ZEMPC_R_F/2Layer/u_opt_Richards_0.02pn.txt")
y_ODE = np.loadtxt("ZEMPC_R_F/2Layer/y_opt_Richards_0.02pn.txt") 

u_LSTM_2p = np.loadtxt("ZEMPC_R_F/LSTM/u_opt_018_023_LSTM_0.02pn.txt")
y_LSTM_2p = np.loadtxt("ZEMPC_R_F/LSTM/y_opt_018_023_LSTM_0.02pn.txt") 

u_2L_2p = np.loadtxt("ZEMPC_R_F/2Layer/u_opt_018_023_0.02pn.txt")
y_2L_2p = np.loadtxt("ZEMPC_R_F/2Layer/y_opt_018_023_0.02pn.txt") 

u_LSTM_2m = np.loadtxt("ZEMPC_R_F/LSTM/u_opt_018_023_LSTM_0.02mn.txt")
y_LSTM_2m = np.loadtxt("ZEMPC_R_F/LSTM/y_opt_018_023_LSTM_0.02mn.txt") 

u_2L_2m = np.loadtxt("ZEMPC_R_F/2Layer/u_opt_018_023_0.02mn.txt")
y_2L_2m = np.loadtxt("ZEMPC_R_F/2Layer/y_opt_018_023_0.02mn.txt") 

u_LSTM_5p = np.loadtxt("ZEMPC_R_F/LSTM/u_opt_018_023_LSTM_0.05pn.txt")
y_LSTM_5p = np.loadtxt("ZEMPC_R_F/LSTM/y_opt_018_023_LSTM_0.05pn.txt") 

u_2L_5p = np.loadtxt("ZEMPC_R_F/2Layer/u_opt_018_023_0.05pn.txt")
y_2L_5p = np.loadtxt("ZEMPC_R_F/2Layer/y_opt_018_023_0.05pn.txt") 

u_2L_5m = np.loadtxt("ZEMPC_R_F/2Layer/u_opt_018_023_0.05mn.txt")
y_2L_5m = np.loadtxt("ZEMPC_R_F/2Layer/y_opt_018_023_0.05mn.txt") 

u_LSTM_5m = np.loadtxt("ZEMPC_R_F/LSTM/u_opt_018_023_LSTM_0.05mn.txt")
y_LSTM_5m = np.loadtxt("ZEMPC_R_F/LSTM/y_opt_018_023_LSTM_0.05mn.txt") 


u_LSTM_2p_2n = np.loadtxt("ZEMPC_R_F/LSTM/u_opt_018_023_LSTM_0.02pn_0.02mn.txt")
y_LSTM_2p_2n = np.loadtxt("ZEMPC_R_F/LSTM/y_opt_018_023_LSTM_0.02pn_0.02mn.txt") 

u_2L_2p_2n = np.loadtxt("ZEMPC_R_F/2Layer/u_opt_018_023_0.02pn_0.02mn.txt")
y_2L_2p_2n = np.loadtxt("ZEMPC_R_F/2Layer/y_opt_018_023_0.02pn_0.02mn.txt") 


u_LSTM_5p_5n = np.loadtxt("ZEMPC_R_F/LSTM/u_opt_018_023_LSTM_0.05pn_0.05mn.txt")
y_LSTM_5p_5n = np.loadtxt("ZEMPC_R_F/LSTM/y_opt_018_023_LSTM_0.05pn_0.05mn.txt") 

u_2L_5p_5n = np.loadtxt("ZEMPC_R_F/2Layer/u_opt_018_023_0.05pn_0.05mn.txt")
y_2L_5p_5n = np.loadtxt("ZEMPC_R_F/2Layer/y_opt_018_023_0.05pn_0.05mn.txt") 


y_lb = np.ones(u_2L_2p.size)*0.18
y_ub = np.ones(y_2L_2p.size)*0.23




lbZ = 0.18
ubZ = 0.23
Q = 4000
R = 100
ubu = 0
lbu = -1e-06
def obj(u, y, Q, R):
    yz = max(lbZ, y)
    J = mtimes((y - yz).T, mtimes(Q, (y - yz))) + R*(1 - u)**2
    return J

obj_LSTM_2p = 0
obj_2L_2p = 0
obj_LSTM_5p = 0
obj_2L_5p = 0
obj_LSTM_2p_2n = 0
obj_2L_2p_2n = 0
obj_LSTM_5p_5n = 0
obj_2L_5p_5n = 0
for i in range(u_2L_2p.size):
    uLs_2p = (u_LSTM_2p[i] - lbu)/(ubu-lbu)
    u2Ls_2p = (u_2L_2p[i] - lbu)/(ubu-lbu)
    uLs_5p = (u_LSTM_5p[i] - lbu)/(ubu-lbu)
    u2Ls_5p = (u_2L_5p[i] - lbu)/(ubu-lbu)
    uLs_2p_2n = (u_LSTM_2p_2n[i] - lbu)/(ubu-lbu)
    u2Ls_2p_2n = (u_2L_2p_2n[i] - lbu)/(ubu-lbu)
    uLs_5p_5n = (u_LSTM_5p_5n[i] - lbu)/(ubu-lbu)
    u2Ls_5p_5n = (u_2L_5p_5n[i] - lbu)/(ubu-lbu)
    
    
    obj_LSTM_2p += obj(uLs_2p, y_LSTM_2p[i+1], Q, R)    
    obj_2L_2p += obj(u2Ls_2p, y_2L_2p[i+1], Q, R)   
    obj_LSTM_5p += obj(uLs_5p, y_LSTM_5p[i+1], Q, R)    
    obj_2L_5p += obj(u2Ls_5p, y_2L_5p[i+1], Q, R) 
    obj_LSTM_2p_2n += obj(uLs_2p_2n, y_LSTM_2p_2n[i+1], Q, R)    
    obj_2L_2p_2n += obj(u2Ls_2p_2n, y_2L_2p_2n[i+1], Q, R) 
    obj_LSTM_5p_5n += obj(uLs_5p_5n, y_LSTM_5p_5n[i+1], Q, R)    
    obj_2L_5p_5n += obj(u2Ls_5p_5n, y_2L_5p_5n[i+1], Q, R) 
    
   
# print(np.array(obj_LSTM_2p))
# print(np.array(obj_2L_2p))
# print(np.array(obj_LSTM_5p))
# print(np.array(obj_2L_5p))
# print(np.array(obj_LSTM_2p_2n))
# print(np.array(obj_2L_2p_2n))
# print(np.array(obj_LSTM_5p_5n))
# print(np.array(obj_2L_5p_5n))

# %%
tplot = np.linspace(0, Nsim/3600, num = int(Nsim/Delta), endpoint=False)
csfont = {'fontname':'Times New Roman'}
t_end = 59
fig, axs = plt.subplots(2,sharex = True, figsize=(11,7))
font = font_manager.FontProperties(family='Times New Roman',
                                   # weight='bold',
                                    style='normal', size=19)

axs[0].step(tplot[:t_end], -u_ODE[:t_end], marker = "*", color = "b", label = "Richards Eqn.")
axs[0].step(tplot[:t_end], -u_LSTM_2p[:t_end], marker = ".", color = "m", label = "Single LSTM")
axs[0].step(tplot[:t_end], -u_2L_2p[:t_end], marker = "v", color = "y", label = "2-layered NN")
# plt.legend(frameon=False, ncol=1, loc="upper right", prop=font_l)
axs[0].set_ylabel('Irrigation Rate ($m/s$)', labelpad=30, fontproperties=font)
axs[1].set_xlabel("Time (hr)", fontproperties=font)
axs[1].set_ylabel('Soil Moisture ($m^3/m^3$)', labelpad=5, fontproperties=font)


# plt.figure(2)
axs[1].plot(tplot[:t_end+1], y_lb[:t_end+1], linestyle = "-.", color = "k", linewidth = 2)
axs[1].plot(tplot[:t_end+1], y_ub[:t_end+1], linestyle = "-.", color = "k", linewidth = 2)
axs[1].plot(tplot[:t_end+1], y_ODE[:t_end+1], marker = "*", color = "b", label = "Richards Eqn.")
axs[1].plot(tplot[:t_end+1], y_LSTM_2p[:t_end+1], marker = ".", color = "m", label = "Single LSTM")
axs[1].plot(tplot[:t_end+1], y_2L_2p[:t_end+1], marker = "v", color = "y", label = "2-layer NN")

axs[1].legend(frameon=False, ncol=1, loc="upper right", prop=font)
axs[0].tick_params(axis='y', labelsize=15)
axs[1].tick_params(axis='y', labelsize=15)
axs[1].tick_params(axis='x', labelsize=15)
plt.show()
# %% 5% Process Noise - ODE, LSTM & 2-layered
# u_ODE = np.loadtxt("ZEMPC/Noise/u_opt_Richards_0.05pn.txt")
# y_ODE = np.loadtxt("ZEMPC/Noise/y_opt_Richards_0.05pn.txt") 





y_lb = np.ones(u_2L_2p.size)*0.18
y_ub = np.ones(y_2L_2p.size)*0.23

tplot = np.linspace(0, Nsim/3600, num = int(Nsim/Delta), endpoint=False)

t_end = 59
fig, axs = plt.subplots(2,sharex = True, figsize=(11,7))
axs[0].step(tplot[:t_end], -u_LSTM_5p_5n[:t_end], color = "m", marker = ".", label = "Single LSTM")
axs[0].step(tplot[:t_end], -u_2L_5p_5n[:t_end], color = "y", marker = "v", label = "2-layer NN")
# axs[1].plot(t[:a],y[:a]) #, label = 'Product Concentration')
axs[0].set_ylabel('Irrigation Rate ($m/s$)', labelpad=15, fontproperties=font)
axs[1].set_xlabel("Time (hr)", fontproperties=font)
axs[1].set_ylabel('Soil Moisture ($m^3/m^3$)', labelpad=5, fontproperties=font)
# for ax in axs.flatten():
#     labels = ax.get_xticklabels() + ax.get_yticklabels()
#     [label.set_fontname('Times New Roman') for label in labels]
#     [label.set_fontsize('xx-large') for label in labels]
    

csfont = {'fontname':'Times New Roman'}
# plt.figure(2,figsize=(6, 4.5))
# plt.style.use('grayscale')
axs[1].plot(tplot[:t_end+1], y_lb[:t_end+1], linestyle = "-.", color = "k", linewidth = 2)
axs[1].plot(tplot[:t_end+1], y_ub[:t_end+1], linestyle = "-.", color = "k", linewidth = 2)
axs[1].plot(tplot[:t_end+1], y_LSTM_5p_5n[:t_end+1], color = "m", marker = ".", label = "Single LSTM")
axs[1].plot(tplot[:t_end+1], y_2L_5p_5n[:t_end+1], color = "y", marker = "v", label = "2-layer NN")
# plt.xlabel('Time (min)', **csfont, fontsize = 17)
# plt.ylabel('Product concentration ($mole/m^3$)', **csfont, fontsize = 17)
axs[1].legend(frameon=False, ncol=1, loc="upper right", prop=font)
axs[0].tick_params(axis='y', labelsize=15)
axs[1].tick_params(axis='y', labelsize=15)
axs[1].tick_params(axis='x', labelsize=15)
plt.show()


# plt.figure(3)
# # plt.step(tplot[:t_end], -u_ODE[:t_end], linestyle = "dashed", color = "b", label = "Richards Eqn.")
# # plt.step(tplot[:t_end], -u_LSTM_2p[:t_end], linestyle = "dotted", color = "m", linewidth = 2, label = "Single LSTM - 2% P.N.")
# # plt.step(tplot[:t_end], -u_2L_2p[:t_end], linestyle = "solid", color = "y", label = "2-layered NN - 2% P.N.")
# # plt.step(tplot[:t_end], -u_LSTM_5p[:t_end], color = "m", marker = ".", label = "Single LSTM - 5% P.N.")
# # plt.step(tplot[:t_end], -u_2L_5p[:t_end], color = "y", marker = "v", label = "2-layered NN - 5% P.N.")
# plt.step(tplot[:t_end], -u_LSTM_2p_2n[:t_end], color = "m", marker = ".", label = "Single LSTM - 5% P.N. & M.N.")
# plt.step(tplot[:t_end], -u_2L_2p_2n[:t_end], color = "y", marker = "v", label = "2-layered NN - 5% P.N. & M.N.")
# plt.legend(frameon=False, ncol=1, loc="upper right", prop=font_l)
# plt.xlabel('Time (hr)', **csfont, fontsize = 17)
# plt.ylabel('Irrigation (m/s)', **csfont, fontsize = 17)


# plt.figure(4)
# plt.plot(tplot[:t_end+1], y_lb[:t_end+1], linestyle = "-.", color = "k", linewidth = 2)
# plt.plot(tplot[:t_end+1], y_ub[:t_end+1], linestyle = "-.", color = "k", linewidth = 2)
# # plt.plot(tplot[:t_end+1], y_ODE[:t_end+1], linestyle = "dashed", color = "b", label = "Richards Eqn.")
# # plt.plot(tplot[:t_end+1], y_LSTM_2p[:t_end+1], linestyle = "dotted", color = "m", label = "Single LSTM - 2% P.N.")
# # plt.plot(tplot[:t_end+1], y_2L_2p[:t_end+1], linestyle = "solid", color = "y", label = "2-layered NN - 2% P.N.")
# # plt.plot(tplot[:t_end+1], y_LSTM_5p[:t_end+1], linestyle = "dashdot", color = "m", label = "Single LSTM - 5% P.N.")
# # plt.plot(tplot[:t_end+1], y_2L_5p[:t_end+1], linestyle = "solid", color = "m", label = "2-layered NN - 5% P.N.")
# plt.plot(tplot[:t_end+1], y_LSTM_2p_2n[:t_end+1], color = "m", marker = ".", label = "Single LSTM - 5% P.N. M.N.")
# plt.plot(tplot[:t_end+1], y_2L_2p_2n[:t_end+1], color = "y", marker = "v", label = "2-layered NN - 5% P.N. & M.N.")
# plt.legend(frameon=False, ncol=1, loc="lower right", prop=font_l)
# plt.xlabel('Time (hr)', **csfont, fontsize = 17)
# plt.ylabel('Soil Moisture ($m^3/m^3$)', **csfont, fontsize = 17)

# plt.figure(5)
# # plt.step(tplot[:t_end], -u_ODE[:t_end], linestyle = "dashed", color = "b", label = "Richards Eqn.")
# # plt.step(tplot[:t_end], -u_LSTM_2p[:t_end], linestyle = "dotted", color = "m", linewidth = 2, label = "Single LSTM - 2% P.N.")
# # plt.step(tplot[:t_end], -u_2L_2p[:t_end], linestyle = "solid", color = "y", label = "2-layered NN - 2% P.N.")
# plt.step(tplot[:t_end], -u_LSTM_5p[:t_end], color = "m", marker = ".", label = "Single LSTM - 5% P.N.")
# plt.step(tplot[:t_end], -u_2L_5p[:t_end], color = "y", marker = "v", label = "2-layered NN - 5% P.N.")
# # plt.step(tplot[:t_end], -u_LSTM_5p_5n[:t_end], color = "m", marker = ".", label = "Single LSTM - 5% P.N. & M.N.")
# # plt.step(tplot[:t_end], -u_2L_5p_5n[:t_end], color = "y", marker = "v", label = "2-layered NN - 5% P.N. & M.N.")
# plt.legend(frameon=False, ncol=1, loc="upper right", prop=font_l)
# plt.xlabel('Time (hr)', **csfont, fontsize = 17)
# plt.ylabel('Irrigation (m/s)', **csfont, fontsize = 17)


# plt.figure(6)
# plt.plot(tplot[:t_end+1], y_lb[:t_end+1], linestyle = "-.", color = "k", linewidth = 2)
# plt.plot(tplot[:t_end+1], y_ub[:t_end+1], linestyle = "-.", color = "k", linewidth = 2)
# # plt.plot(tplot[:t_end+1], y_ODE[:t_end+1], linestyle = "dashed", color = "b", label = "Richards Eqn.")
# # plt.plot(tplot[:t_end+1], y_LSTM_2p[:t_end+1], linestyle = "dotted", color = "m", label = "Single LSTM - 2% P.N.")
# # plt.plot(tplot[:t_end+1], y_2L_2p[:t_end+1], linestyle = "solid", color = "y", label = "2-layered NN - 2% P.N.")
# plt.plot(tplot[:t_end+1], y_LSTM_5p[:t_end+1], marker = ".", color = "m", label = "Single LSTM - 5% P.N.")
# plt.plot(tplot[:t_end+1], y_2L_5p[:t_end+1], marker = ".", color = "y", label = "2-layered NN - 5% P.N.")
# # plt.plot(tplot[:t_end+1], y_LSTM_5p_5n[:t_end+1], color = "m", marker = ".", label = "Single LSTM - 5% P.N. M.N.")
# # plt.plot(tplot[:t_end+1], y_2L_5p_5n[:t_end+1], color = "y", marker = "v", label = "2-layered NN - 5% P.N. & M.N.")
# plt.legend(frameon=False, ncol=1, loc="lower right", prop=font_l)
# plt.xlabel('Time (hr)', **csfont, fontsize = 17)
# plt.ylabel('Soil Moisture ($m^3/m^3$)', **csfont, fontsize = 17)
# lbZ = 0.18
# ubZ = 0.23
# Q = 4000
# R = 100
# def obj(u, y, Q, R):
#     yz = max(lbZ, y)
#     J = mtimes((y - yz).T, mtimes(Q, (y - yz))) + R*(1 - u)**2
#     return J

# obj_LSTM = 0
# obj_2L = 0

# for i in range(u_2L.size):
#     obj_LSTM += obj(u_LSTM[i], y_LSTM[i+1], Q, R)    
#     obj_2L += obj(u_2L[i], y_2L[i+1], Q, R)   
    
# print(np.array(obj_LSTM))
# print(np.array(obj_2L))


# %% 5% pn + 5% mn - ODE, LSTM & 2-layered




y_lb = np.ones(u_2L.size+1)*0.18
y_ub = np.ones(y_2L.size+1)*0.23

tplot = np.linspace(0, Nsim/3600, num = int(Nsim/Delta), endpoint=False)

t_end = 40
plt.figure(3)
# plt.step(tplot[:t_end], -u_ODE[:t_end], linestyle = "dashed", color = "b", label = "Richards Eqn.")
plt.step(tplot[:t_end], -u_LSTM_5p_5n[:t_end], linestyle = "dotted", color = "m", linewidth = 2, label = "Single LSTM - 5% P.N. & M.N.")
plt.step(tplot[:t_end], -u_2L_5p_5n[:t_end], linestyle = "solid", color = "y", label = "2-layered NN - 5% P.N. & M.N.")
plt.legend(frameon=False, ncol=1, loc="upper right", prop=font_l)
plt.xlabel('Time (hr)', **csfont, fontsize = 17)
plt.ylabel('Irrigation (m/s)', **csfont, fontsize = 17)


plt.figure(4)
# plt.plot(tplot[:t_end+1], y_lb[:t_end+1], linestyle = "-.", color = "k", linewidth = 2)
# plt.plot(tplot[:t_end+1], y_ub[:t_end+1], linestyle = "-.", color = "k", linewidth = 2)
# plt.plot(tplot[:t_end+1], y_ODE[:t_end+1], linestyle = "dashed", color = "b", label = "Richards Eqn.")
plt.plot(tplot[:t_end+1], y_LSTM_5p_5n[:t_end+1], linestyle = "dotted", color = "m", label = "Single LSTM - 5% P.N. & M.N.")
plt.plot(tplot[:t_end+1], y_2L_5p_5n[:t_end+1], linestyle = "solid", color = "y", label = "2-layered NN - 5% P.N. & M.N.")
plt.legend(frameon=False, ncol=1, loc="lower right", prop=font_l)
plt.xlabel('Time (hr)', **csfont, fontsize = 17)
plt.ylabel('Soil Moisture ($m^3/m^3$)', **csfont, fontsize = 17)


lbZ = 0.18
ubZ = 0.23
Q = 4000
R = 100
def obj(u, y, Q, R):
    yz = max(lbZ, y)
    J = mtimes((y - yz).T, mtimes(Q, (y - yz))) + R*(1 - u)**2
    return J

obj_LSTM = 0
obj_2L = 0

for i in range(u_2L.size):
    obj_LSTM += obj(u_LSTM[i], y_LSTM[i+1], Q, R)    
    obj_2L += obj(u_2L[i], y_2L[i+1], Q, R)   
    
print(np.array(obj_LSTM))
print(np.array(obj_2L))


# %% Asym Zone - Rain
u_20_5 = np.loadtxt("ZEMPC_R_F/AsymZone/u_opt_rain_20_5_asymZ_Ql6000_Qu4000_0.05mn.txt")
y_20_5 = np.loadtxt("ZEMPC_R_F/AsymZone/y_opt_rain_20_5_asymZ_Ql6000_Qu4000_0.05mn.txt") 

u_30_10_84 = np.loadtxt("ZEMPC_R_F/AsymZone/u_opt_rain_30_10_asymZ_Ql8000_Qu4000_0.05mn.txt")
y_30_10_84 = np.loadtxt("ZEMPC_R_F/AsymZone/y_opt_rain_30_10_asymZ_Ql8000_Qu4000_0.05mn.txt") 

u_30_10_64 = np.loadtxt("ZEMPC_R_F/AsymZone/u_opt_rain_30_10_asymZ_Ql6000_Qu4000_0.05mn.txt")
y_30_10_64 = np.loadtxt("ZEMPC_R_F/AsymZone/y_opt_rain_30_10_asymZ_Ql6000_Qu4000_0.05mn.txt") 

u_30_10_42 = np.loadtxt("ZEMPC_R_F/AsymZone/u_opt_rain_30_10_asymZ_Ql4000_Qu2000_0.05mn.txt")
y_30_10_42 = np.loadtxt("ZEMPC_R_F/AsymZone/y_opt_rain_30_10_asymZ_Ql4000_Qu2000_0.05mn.txt") 

u_30_10 = np.loadtxt("ZEMPC_R_F/AsymZone/u_opt_rain_30_10_symZ_Q4000_0.05mn.txt")
y_30_10 = np.loadtxt("ZEMPC_R_F/AsymZone/y_opt_rain_30_10_symZ_Q4000_0.05mn.txt") 

# u_40_10 = np.loadtxt("ZEMPC_R_F/AsymZone/u_opt_rain_40_10_asymZ_Ql6000_Qu4000_0.05mn.txt")
# y_40_10 = np.loadtxt("ZEMPC_R_F/AsymZone/y_opt_rain_40_10_asymZ_Ql6000_Qu4000_0.05mn.txt") 

# us_40_10 = np.loadtxt("ZEMPC_R_F/AsymZone/u_opt_rain_40_10_Ql4000_Qu4000.txt")
# ys_40_10 = np.loadtxt("ZEMPC_R_F/AsymZone/y_opt_rain_40_10_Ql4000_Qu4000.txt") 

# u_50_10 = np.loadtxt("ZEMPC_R_F/AsymZone/u_opt_rain_50_10_asymZ_Ql4000_Qu2000.txt")
# y_50_10 = np.loadtxt("ZEMPC_R_F/AsymZone/y_opt_rain_50_10_asymZ_Ql4000_Qu2000.txt") 

# us_50_10 = np.loadtxt("ZEMPC_R_F/AsymZone/u_opt_rain_50_10_Ql4000_Qu4000.txt")
# ys_50_10 = np.loadtxt("ZEMPC_R_F/AsymZone/y_opt_rain_50_10_Ql4000_Qu4000.txt") 

# y_lb = np.ones(u_2L.size)*0.18
# y_ub = np.ones(y_2L.size)*0.23

tplot = np.linspace(0, Nsim/3600, num = int(Nsim/Delta), endpoint=False)

t_end = 59
plt.figure(7)
# plt.step(tplot[:t_end], -u_ODE[:t_end], linestyle = "dashed", color = "b", label = "Richards Eqn.")
# plt.step(tplot[:t_end], -us_40_10[:t_end], linestyle = "dotted", color = "m", linewidth = 2, label = "Sym. Zone")

plt.step(tplot[:t_end], -u_30_10[:t_end], linestyle = "solid", color = "b", label = "Sym. Zone")
plt.step(tplot[:t_end], -u_30_10_42[:t_end], linestyle = "-.", color = "m", label = "Asym. Zone-Low")
plt.step(tplot[:t_end], -u_30_10_84[:t_end], linestyle = "--", color = "y", label = "Asym. Zone-High")
# plt.step(tplot[:t_end], -u_20_5[:t_end], linestyle = "solid", color = "y", label = "20_5")
plt.legend(frameon=False, ncol=1, loc="upper right", prop=font_l)
plt.xlabel('Time (hr)', **csfont, fontsize = 17)
plt.ylabel('Irrigation (m/s)', **csfont, fontsize = 17)


plt.figure(8)
plt.plot(tplot[:t_end+1], y_lb[:t_end+1], linestyle = "-.", color = "k", linewidth = 2)
plt.plot(tplot[:t_end+1], y_ub[:t_end+1], linestyle = "-.", color = "k", linewidth = 2)
# plt.plot(tplot[:t_end+1], y_ODE[:t_end+1], linestyle = "dashed", color = "b", label = "Richards Eqn.")
# plt.plot(tplot[:t_end+1], y_40_10[:t_end+1], linestyle = "solid", color = "m", label = "40_10")
plt.plot(tplot[:t_end+1], y_30_10[:t_end+1], linestyle = "solid", color = "b", label = "30_10")
plt.plot(tplot[:t_end+1], y_30_10_42[:t_end+1], linestyle = "-.", color = "m", label = "Asym. Zone-Low")
plt.plot(tplot[:t_end+1], y_30_10_84[:t_end+1], linestyle = "--", color = "y", label = "Asym. Zone-High")
# plt.plot(tplot[:t_end+1], y_40_10[:t_end+1], linestyle = "solid", color = "y", label = "Asym. Zone")
plt.legend(frameon=False, ncol=1, loc="lower right", prop=font_l)
plt.xlabel('Time (hr)', **csfont, fontsize = 17)
plt.ylabel('Soil Moisture ($m^3/m^3$)', **csfont, fontsize = 17)

# %% Shrinking Zone

random.seed(5)
np.random.seed(5)
refEvap = np.random.normal(3, 0.2, 10)/1e3/(24*60*60)
Evap = np.zeros(int(Nsim/Delta))
for i in range(10):
    Evap[i*12:i*12+12] = refEvap[i]
# np.savetxt("C:/Users/zhiyinan/OneDrive - ualberta.ca/Irrigation/crop_model_1D/ZMPC Simulation Results/Evap.txt", Evap)
# data_hr = pd.read_csv('C:/Users/zhiyinan/OneDrive - ualberta.ca/Irrigation/crop_model_1D\Rain_hourly.csv')
# precip = data_hr.iloc[240:480,2].values/1e3/(60*60*24)

Precip = np.zeros((5,12))

for i in range(5):
    start = random.randrange(5)
    end = random.randrange(10)
    for j in range(start, end):
        Precip[i,j] = -(random.random()*(20-5)+5)/1e3/(24*60*60)
        # Precip_p[i,j] = Precip[i,j]  + np.random.choice([-1,1]) * 0.3*np.random.random_sample()*Precip[i,j]

Precip = Precip.flatten()
Precip_p = np.zeros(Precip.shape[0])
for i in range(Precip.shape[0]):
        Precip_p[i] = min(0, Precip[i]  + np.random.choice([-1,1])* 0.2*np.random.random_sample()*min(Precip))

Precip_p = np.concatenate((Precip_p.flatten(), np.zeros(20)))
# Precip_p = -abs((Precip*np.random.normal(1,0.8,size = (Precip.shape)))) # uniform

# Precip = np.zeros(int(Nsim/Delta))

# for i in range(int(Nsim/Delta)):
#     Precip[i] = random.choice([0,0,0,0,1]) * (random.random()*(2-0.2)+0.2)/1e3/(24*60*60)

plt.figure(3)
plt.step(np.arange(-Precip.shape[0]), Precip)
plt.step(np.arange(-Precip.shape[0]), Precip_p[:60])

# %%
y_lb_L = np.zeros(int(Nt/Delta)+1)
y_ub_L = np.zeros(int(Nt/Delta)+1)
for i in range(int(Nt/Delta)+1):
    y_lb_L[i] = min(0.18*exp(i/int(Nt/Delta)), 0.205)
    y_ub_L[i] = max(0.23*exp(-i/int(Nt/Delta)), 0.205)

y_lb_7 = np.zeros(int(Nt/Delta)+1)
y_ub_7 = np.zeros(int(Nt/Delta)+1)
for i in range(int(Nt/Delta)+1):
    y_lb_7[i] = min(0.18*exp(i/int(Nt/Delta)*0.7), 0.20)
    y_ub_7[i] = max(0.23*exp(-i/int(Nt/Delta)*0.7), 0.21)
    
y_lb_5 = np.zeros(int(Nt/Delta)+1)
y_ub_5 = np.zeros(int(Nt/Delta)+1)
for i in range(int(Nt/Delta)+1):
    y_lb_5[i] = min(0.18*exp(i/int(Nt/Delta)*0.5), 0.20)
    y_ub_5[i] = max(0.23*exp(-i/int(Nt/Delta)*0.5), 0.21)
    
y_lb_B = np.zeros(int(Nt/Delta)+1)
y_ub_B = np.zeros(int(Nt/Delta)+1)
for i in range(int(Nt/Delta)+1):
    y_lb_B[i] = min(0.18*exp(i/int(Nt/Delta)), 0.20)
    y_ub_B[i] = max(0.23*exp(-i/int(Nt/Delta)), 0.21)
    
ax1 = plt.figure(2)
plt.plot(np.arange(0,42,2),y_lb_L, linestyle = "--", color = "m", label = r"$\underbar \!\!\! Y = \overline{Y}, \mu = 1$")
plt.plot(np.arange(0,42,2),y_ub_L, linestyle = "--", color = "m")
plt.fill_between(np.arange(0,42,2), y_lb_L, y_ub_L, alpha=0.2, color = "b")
plt.plot(np.arange(0,42,2),y_lb_B, linestyle = "-.", color = "b", label = r"$\underbar \!\!\! Y < \overline{Y}, \mu = 1$")
plt.plot(np.arange(0,42,2),y_ub_B, linestyle = "-.", color = "b")
plt.fill_between(np.arange(0,42,2), y_lb_B, y_ub_B, alpha=0.2)
plt.legend(frameon=False, ncol=1, loc="upper right", prop=font_l)
plt.xlabel('Time (hr)', **csfont, fontsize = 17)
plt.ylabel('Soil Moisture ($m^3/m^3$)', **csfont, fontsize = 17)
ax1.set_rasterized(True)
plt.tight_layout()
plt.savefig('Shrk_Z_B_L.eps')


ax = plt.figure(3)
plt.plot(np.arange(0,42,2),y_lb_B, linestyle = "-.", color = "b", label = r"$\underbar \!\!\! Y < \overline{Y}, \mu = 1$")
plt.plot(np.arange(0,42,2),y_ub_B, linestyle = "-.", color = "b")
plt.fill_between(np.arange(0,42,2), y_lb_B, y_ub_B, alpha=0.2)
plt.plot(np.arange(0,42,2),y_lb_7, linestyle = ":", color = "r", label = r"$\underbar \!\!\! Y < \overline{Y}, \mu = 0.7$")
plt.plot(np.arange(0,42,2),y_ub_7, linestyle = ":", color = "r")
plt.fill_between(np.arange(0,42,2), y_lb_7, y_ub_7, alpha=0.2)
plt.plot(np.arange(0,42,2),y_lb_5, linestyle = "-", color = "c", label = r"$\underbar \!\!\! Y < \overline{Y}, \mu = 0.5$")
plt.plot(np.arange(0,42,2),y_ub_5, linestyle = "-", color = "c")
plt.fill_between(np.arange(0,42,2), y_lb_5, y_ub_5, alpha=0.2)#, color = "c"
plt.legend(frameon=False, ncol=1, loc="upper right", prop=font_l)
plt.xlabel('Time (hr)', **csfont, fontsize = 17)
plt.ylabel('Soil Moisture ($m^3/m^3$)', **csfont, fontsize = 17)
ax.set_rasterized(True)
plt.tight_layout()
plt.savefig('Shrk_Z_B.eps')

# %%
u_20_5_sz_B = np.loadtxt("ZEMPC_R_F/ShrinkingZone/u_opt_018_023_rain_20_5_0.05mn_shrinkZ_B.txt")
y_20_5_sz_B = np.loadtxt("ZEMPC_R_F/ShrinkingZone/y_opt_018_023_rain_20_5_0.05mn_shrinkZ_B.txt") 

u_20_5_sz_B_5 = np.loadtxt("ZEMPC_R_F/ShrinkingZone/u_opt_018_023_rain_20_5_0.05mn_shrinkZ_B_0.5.txt")
y_20_5_sz_B_5 = np.loadtxt("ZEMPC_R_F/ShrinkingZone/y_opt_018_023_rain_20_5_0.05mn_shrinkZ_B_0.5.txt") 

u_20_5_sz_B_7 = np.loadtxt("ZEMPC_R_F/ShrinkingZone/u_opt_018_023_rain_20_5_0.05mn_shrinkZ_B_0.7.txt")
y_20_5_sz_B_7 = np.loadtxt("ZEMPC_R_F/ShrinkingZone/y_opt_018_023_rain_20_5_0.05mn_shrinkZ_B_0.7.txt") 
# u_30_10_sz_B = np.loadtxt("ZEMPC_R_F/ShrinkingZone/u_opt_018_023_rain_30_10_0.05mn_shrinkZ_B.txt")
# y_30_10_sz_B = np.loadtxt("ZEMPC_R_F/ShrinkingZone/y_opt_018_023_rain_30_10_0.05mn_shrinkZ_B.txt") 

u_20_5_sz_L = np.loadtxt("ZEMPC_R_F/ShrinkingZone/u_opt_018_023_rain_20_5_0.05mn_shrinkZ_L.txt")
y_20_5_sz_L = np.loadtxt("ZEMPC_R_F/ShrinkingZone/y_opt_018_023_rain_20_5_0.05mn_shrinkZ_L.txt") 

# u_30_10_sz_L = np.loadtxt("ZEMPC_R_F/ShrinkingZone/u_opt_018_023_rain_30_10_0.05mn_shrinkZ_L.txt")
# y_30_10_sz_L = np.loadtxt("ZEMPC_R_F/ShrinkingZone/y_opt_018_023_rain_30_10_0.05mn_shrinkZ_L.txt") 
font = font_manager.FontProperties(family='Times New Roman',
                                   # weight='bold',
                                    style='normal', size=19)

y_lb = np.ones(61)*0.18
y_ub = np.ones(61)*0.23

tplot = np.linspace(0, Nsim/3600, num = int(Nsim/Delta)+1, endpoint=True)
t_end = 60
fig, axs = plt.subplots(3,sharex = True, figsize=(11,15))
axs[0].step(tplot[:t_end], -Precip, color = "k", label = "Act. Precip.")
axs[0].step(tplot[:t_end], -Precip_p[:60], color = "r", linestyle = "--", label = "Pred. Precip.")
axs[0].legend(frameon=False, ncol=1, loc="upper right", prop=font)
# plt.step(tplot[:t_end], -u_ODE[:t_end], linestyle = "dashed", color = "b", label = "Richards Eqn.")
# plt.step(tplot[:t_end], -us_40_10[:t_end], linestyle = "dotted", color = "m", linewidth = 2, label = "Sym. Zone")
axs[1].step(tplot[:t_end], -u_20_5[:t_end], linestyle = "solid", color = "r", label = "Time Inv. Z.")

axs[1].step(tplot[:t_end], -u_20_5_sz_L[:t_end], linestyle = "--", color = "m", label = r"$\underbar \!\!\! Y  = \overline{Y}, \mu = 1$")
axs[1].step(tplot[:t_end], -u_20_5_sz_B[:t_end], linestyle = "-.", color = "b", label = r"$\underbar \!\!\! Y  < \overline{Y}, \mu = 1$") # "Zone Term." r"$ \mu$ = 1"
# axs[1].step(tplot[:t_end], -u_20_5_sz_B_7[:t_end], linestyle = "--", color = "c", label = r"$ \mu$ = 0.7")
# axs[1].step(tplot[:t_end], -u_20_5_sz_B_5[:t_end], linestyle = ":", color = "r", label = r"$ \mu$ = 0.5")
axs[0].set_ylabel('Precipitation Rate ($m/s$)', labelpad=15, fontproperties=font)
axs[1].set_ylabel('Irrigation Rate ($m/s$)', labelpad=28, fontproperties=font)
axs[2].set_xlabel("Time (hr)", fontproperties=font)
axs[2].set_ylabel('Soil Moisture ($m^3/m^3$)', labelpad=2, fontproperties=font)

axs[2].plot(tplot[:t_end+1], y_lb[:t_end+1], linestyle = "-.", color = "k", linewidth = 2)
axs[2].plot(tplot[:t_end+1], y_ub[:t_end+1], linestyle = "-.", color = "k", linewidth = 2)
axs[2].fill_between(tplot[:t_end+1], y_lb[:t_end+1], y_ub[:t_end+1], alpha=0.2)
axs[2].plot(tplot[:t_end+1], y_20_5[:t_end+1], linestyle = "solid", color = "r", label = "Time Inv. Z.")
axs[2].plot(tplot[:t_end+1], y_20_5_sz_L[:t_end+1], linestyle = "--", color = "m", label = r"$\underbar \!\!\! Y  = \overline{Y}, \mu = 1$")
axs[2].plot(tplot[:t_end+1], y_20_5_sz_B[:t_end+1], linestyle = "-.", color = "b", label = r"$\underbar \!\!\! Y  < \overline{Y}, \mu = 1$")
# axs[2].plot(tplot[:t_end+1], y_20_5_sz_B_7[:t_end+1], linestyle = "--", color = "c", label = r"$ \mu$ = 0.7")
# axs[2].plot(tplot[:t_end+1], y_20_5_sz_B_5[:t_end+1], linestyle = ":", color = "r", label = r"$ \mu$ = 0.5")

font = font_manager.FontProperties(family='Times New Roman',
                                   # weight='bold',
                                    style='normal', size=18)

axs[2].legend(frameon=False, ncol=1, loc="upper right", prop=font)
axs[0].tick_params(axis='y', labelsize=15)
axs[1].tick_params(axis='y', labelsize=15)
axs[2].tick_params(axis='y', labelsize=15)
axs[2].tick_params(axis='x', labelsize=15)
fig.set_rasterized(True)
plt.tight_layout()
plt.show()
plt.savefig('Shrk_ZEMPC_B_L.eps')


# %%
tplot = np.linspace(0, Nsim/3600, num = int(Nsim/Delta)+1, endpoint=True)
t_end = 60
fig, axs = plt.subplots(2,sharex = True, figsize=(11,8))
# axs[0].step(tplot[:t_end], -Precip, label = "Act. Precip.")
# axs[0].step(tplot[:t_end], -Precip_p[:60], label = "Pred. Precip.")
# axs[0].legend(frameon=False, ncol=1, loc="upper right", prop=font)
# plt.step(tplot[:t_end], -u_ODE[:t_end], linestyle = "dashed", color = "b", label = "Richards Eqn.")
# plt.step(tplot[:t_end], -us_40_10[:t_end], linestyle = "dotted", color = "m", linewidth = 2, label = "Sym. Zone")
# axs[0].step(tplot[:t_end], -u_20_5[:t_end], linestyle = "solid", color = "r", label = "Time Inv. Zone")

# axs[0].step(tplot[:t_end], -u_20_5_sz_L[:t_end], linestyle = "--", color = "m", label = "Point Term.")
axs[0].step(tplot[:t_end], -u_20_5_sz_B[:t_end], linestyle = "-.", color = "b", label = r"$ \underbar \!\!\! Y  < \overline{Y}, \mu$ = 1") # "Zone Term." r"$ \mu$ = 1"
axs[0].step(tplot[:t_end], -u_20_5_sz_B_7[:t_end], linestyle = "--", color = "c", label = r"$ \underbar \!\!\! Y  < \overline{Y}, \mu$ = 0.7")
axs[0].step(tplot[:t_end], -u_20_5_sz_B_5[:t_end], linestyle = ":", color = "r", label = r"$\underbar \!\!\! Y  < \overline{Y},  \mu$ = 0.5")
# axs[0].set_ylabel('Precipitation Rate ($m/s$)', labelpad=15, fontproperties=font)
axs[0].set_ylabel('Irrigation Rate ($m/s$)', labelpad=28, fontproperties=font)
axs[1].set_xlabel("Time (hr)", fontproperties=font)
axs[1].set_ylabel('Soil Moisture ($m^3/m^3$)', labelpad=2, fontproperties=font)
 
axs[1].plot(tplot[:t_end+1], y_lb[:t_end+1], linestyle = "-.", color = "k", linewidth = 2)
axs[1].plot(tplot[:t_end+1], y_ub[:t_end+1], linestyle = "-.", color = "k", linewidth = 2)
axs[1].fill_between(tplot[:t_end+1], y_lb[:t_end+1], y_ub[:t_end+1], alpha=0.2)
# axs[1].plot(tplot[:t_end+1], y_20_5[:t_end+1], linestyle = "solid", color = "r", label = "Time Inv. Zone")
# axs[1].plot(tplot[:t_end+1], y_20_5_sz_L[:t_end+1], linestyle = "--", color = "m", label = "Point Term.")
axs[1].plot(tplot[:t_end+1], y_20_5_sz_B[:t_end+1], linestyle = "-.", color = "b", label = r"$ \underbar \!\!\! Y  < \overline{Y}, \mu$ = 1")
axs[1].plot(tplot[:t_end+1], y_20_5_sz_B_7[:t_end+1], linestyle = ":", color = "r", label = r"$ \underbar \!\!\! Y  < \overline{Y}, \mu$ = 0.7")
axs[1].plot(tplot[:t_end+1], y_20_5_sz_B_5[:t_end+1], linestyle = "-", color = "c", label = r"$ \underbar \!\!\! Y  < \overline{Y}, \mu$ = 0.5")

axs[1].legend(frameon=False, ncol=1, loc="upper right", prop=font)
# axs[0].tick_params(axis='y', labelsize=15)
axs[0].tick_params(axis='y', labelsize=15)
axs[1].tick_params(axis='y', labelsize=15)
axs[1].tick_params(axis='x', labelsize=15)
fig.set_rasterized(True)
plt.tight_layout()
plt.show()
plt.savefig('Shrk_ZEMPC_B.eps')
# plt.figure(10)
# # plt.plot(tplot[:t_end+1], y_lb[:t_end+1], linestyle = "-.", color = "k", linewidth = 2)
# # plt.plot(tplot[:t_end+1], y_ub[:t_end+1], linestyle = "-.", color = "k", linewidth = 2)
# # plt.plot(tplot[:t_end+1], y_ODE[:t_end+1], linestyle = "dashed", color = "b", label = "Richards Eqn.")
# plt.plot(tplot[:t_end+1], y_20_5[:t_end+1], linestyle = "solid", color = "r", label = "30_10")
# plt.plot(tplot[:t_end+1], y_20_5_sz_B[:t_end+1], linestyle = "-.", color = "b", label = "Shrk. Zone-1")
# # plt.plot(tplot[:t_end+1], y_20_5_sz_B_7[:t_end+1], linestyle = "solid", color = "b", label = "Shrk. Zone-2")
# # plt.plot(tplot[:t_end+1], y_20_5_sz_B_5[:t_end+1], linestyle = "solid", color = "r", label = "Shrk. Zone-2")
# plt.plot(tplot[:t_end+1], y_20_5_sz_L[:t_end+1], linestyle = "--", color = "m", label = "Shrk. Zone-L")
# # plt.plot(tplot[:t_end+1], y_40_10[:t_end+1], linestyle = "solid", color = "y", label = "Asym. Zone")
# plt.legend(frameon=False, ncol=1, loc="upper right", prop=font_l)
# plt.xlabel('Time (hr)', **csfont, fontsize = 17)
# plt.ylabel('Soil Moisture ($m^3/m^3$)', **csfont, fontsize = 17)




# plt.figure(4)
# plt.plot(tplot, u_rain_12, linestyle = "solid", color = "y", label = "Rain - Updating Linear Correction")
# plt.plot(tplot, u_rain_pn, linestyle = ":", linewidth = 2, color = "b", label = "Rain + P noise - Updating Linear Correction")
# plt.plot(tplot, u_rain_pmn, linestyle = "--", color = "m", label = "Rain + P & M noise - Updating Linear Correction")
# plt.legend(frameon=False, ncol=1, loc="lower right", prop=font_l)
# plt.xlabel('Time Steps', **csfont, fontsize = 17)
# plt.ylabel('Irrigation (m/s)', **csfont, fontsize = 17)

# plt.figure(5)
# plt.plot(tplot, y_lb[:-1], linestyle = "-.", color = "k", linewidth = 3)
# plt.plot(tplot, y_ub[:-1], linestyle = "-.", color = "k", linewidth = 3)
# plt.plot(tplot, y_rain_12[:-1], linestyle = "solid", color = "y", label = "Rain - Updating Linear Correction")
# plt.plot(tplot, y_rain_pn[:-1], linestyle = ":", linewidth = 2, color = "b", label = "Rain + P noise - Updating Linear Correction")
# plt.plot(tplot, y_rain_pmn[:-1], linestyle = "--", color = "m", label = "Rain + P & M noise - Updating Linear Correction")
# plt.legend(frameon=False, ncol=1, loc="lower right", prop=font_l)
# plt.xlabel('Time Steps', **csfont, fontsize = 17)
# plt.ylabel('Soil Moisture ($m^3/m^3$)', **csfont, fontsize = 17)

# plt.figure(6)
# plt.plot(tplot, u_rain_pmn_nc, linestyle = "solid", color = "k", label = "No Correction")
# plt.plot(tplot, u_rain_pmn, linestyle = "--", color = "m", label = "Updating Linear Correction")
# plt.plot(tplot, u_rain_pmn_cb, linestyle = ":", linewidth = 2, color = "b", label = "Constant Correction")
# plt.legend(frameon=False, ncol=1, loc="lower right", prop=font_l)
# plt.xlabel('Time Steps', **csfont, fontsize = 17)
# plt.ylabel('Irrigation (m/s)', **csfont, fontsize = 17)

# plt.figure(7)
# plt.plot(tplot, y_lb[:-1], linestyle = "-.", color = "k", linewidth = 3)
# plt.plot(tplot, y_ub[:-1], linestyle = "-.", color = "k", linewidth = 3)
# plt.plot(tplot, y_rain_pmn_nc[:-1], linestyle = "solid", color = "k", label = "No Correction")
# plt.plot(tplot, y_rain_pmn[:-1], linestyle = "--", color = "m", label = "Updating Linear Correction")
# plt.plot(tplot, y_rain_pmn_cb[:-1], linestyle = ":", linewidth = 2, color = "b", label = "Constant Correction")
# plt.legend(frameon=False, ncol=1, loc="lower right", prop=font_l)
# plt.xlabel('Time Steps', **csfont, fontsize = 17)
# plt.ylabel('Soil Moisture ($m^3/m^3$)', **csfont, fontsize = 17)